import torch 
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import numpy as np

class DatasetM(Dataset):
    def __init__(self, feature, label, seq_len, horizon=1, shift_pre=0):
        self.horizon = horizon
        self.seq_len = seq_len
        self.shift_pre = shift_pre
        self.feature = feature
        self.label = label

    def __getitem__(self,index):   
        label = (self.label)[(index+self.seq_len-self.shift_pre) : (index+self.seq_len+self.horizon-self.shift_pre),:]
        X  = (self.feature)[index : (index+self.seq_len), :,:]
        label = torch.from_numpy(label) #2dim
        X = torch.from_numpy(X) #3dim
        return X.float() ,label.float()

    def __len__(self):
        return len(self.feature)-self.horizon-self.seq_len+self.shift_pre
 

class MyDataset_twoset(Dataset):
    def __init__(self, data_x, data_y,  back_len, horizon, seq_len, times=None, timesenc=False):
        data_y[np.isinf(data_y)] = 0
        data_y[np.isnan(data_y)] = 0
        self.data_x = data_x
        self.data_y = data_y *10
        self.horizon = horizon
        self.seq_len = seq_len
        self.back_len = back_len
        self.timesenc = timesenc
        if timesenc:
            self.times = times
            df_stamp = pd.DataFrame()
            df_stamp['date'] = pd.to_datetime(times)
            df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1)
            df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1)
            df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1)
            df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1)
            df_stamp['min'] = df_stamp.date.apply(lambda row: row.minute, 1)
            data_stamp = df_stamp.drop(columns=['date'])
            data_stamp = data_stamp.values
            self.data_stamp = data_stamp
        
    def __getitem__(self,index):   
        x_begin = index
        x_end = x_begin + self.seq_len
        # y_begin = x_end - self.back_len
        # y_end = y_begin + self.back_len + self.horizon
        seq_x = self.data_x[x_begin:x_end]
        # if y_begin == y_end:
        #     seq_y = np.expand_dims(self.data_y[y_begin],axis=0)
        # else:
        seq_y = np.expand_dims(self.data_y[x_end-1],axis=0)
        if not self.timesenc:
            return seq_x, seq_y
        seq_x_mark = self.data_stamp[x_begin:x_end]
        seq_y_mark = self.data_stamp[x_end-1]
        return seq_x, seq_y, seq_x_mark, seq_y_mark

    def __len__(self):
        return len(self.data_x) - self.seq_len #- self.horizon 